import io
import os
import numpy
import torch
import pandas as pd
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class MyDataset(Dataset):
    def __init__(self, csv_path, txt_path):
        self.txt_path = txt_path
        self.csv_path = csv_path
        self.csvAllData = os.listdir(self.csv_path)
        self.txtAllData = os.listdir(self.txt_path)
        # 将数据转换成tensor形式

    def __len__(self):
        return len(self.csvAllData)

    def __getitem__(self, index):
        label = 1
        txtFileName = os.path.join(self.txt_path, self.txtAllData[index])
        print(txtFileName)
        # print("index:"+str(index))
        csvFileName = os.path.join(self.csv_path, self.csvAllData[index])
        csvData = pd.read_csv(csvFileName, header=None)  # header=None是去掉表头部分
        csvData = csvData.iloc[:, 0]
        csvData = np.array(csvData)
        txtData = numpy.loadtxt(txtFileName, delimiter=',')

        txtData = numpy.reshape(txtData, (4, 1800))

        to_tensor = transforms.Compose([
            transforms.ToTensor()
        ])
        txtData = to_tensor(txtData)
        if csvData[1] - csvData[0] >= 100:
            label = 0
        elif csvData[1] - csvData[0] < 100:
            label = 1
        return label, txtData


myTrainDataset = MyDataset("D:\PycharmProjects\My_project\kneepoint\csv\\train", "D:\PycharmProjects\My_project\kneepoint\data\\train")
print(type(myTrainDataset))
train_loader = torch.utils.data.DataLoader(
    dataset=myTrainDataset,
    batch_size=5,
    shuffle=True,
)

for epoch in range(1000):
    for batch_idx, train_set in enumerate(train_loader):
        data, label = train_set
        print(batch_idx, data.shape, label.shape)
